//
// Created by l00391713 on 2024/3/2.
//
#include <pybind11/pybind11.h>

#include <cstdint>
#include <string>

#include "ParallelTcpStore.hpp"

namespace py = pybind11;

PYBIND11_MODULE(torch_bind, m)
{
    m.doc() = "torch pybind api";
    py::class_<::c10d::ParallelTcpStore>(m, "Parallel")
            .def(py::init([](const std::string &host, uint16_t port, bool isServer){
                ::c10d::TCPStoreOptions options{};
                options.isServer = isServer;
                options.port = port;
                return std::make_unique<::c10d::ParallelTcpStore>(host, options);
            })).def("set", [](::c10d::ParallelTcpStore &store, const std::string &key, const std::string &value) {
                store.set(key, std::vector<uint8_t>(value.begin(), value.end()));
            }).def("add", &::c10d::ParallelTcpStore::add)
            .def("delete", &::c10d::ParallelTcpStore::deleteKey)
            .def("get", [](::c10d::ParallelTcpStore &store, const std::string &key) {
                auto result = store.get(key);
                return std::string(result.begin(), result.end());
            });
}